# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of Cybertron package.
#
# The Cybertron is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
SchNet Model
"""

import torch
from torch import nn, Tensor
from typing import Union, List, Optional, Tuple

from .model import MolecularGNN, _model_register
from ..interaction import Interaction, SchNetInteraction
from ..utils import get_arguments, get_integer

@_model_register('schnet')
class SchNet(MolecularGNN):
    r"""SchNet Model.

    Reference:

        Schütt, K. T.; Sauceda, H. E.; Kindermans, P.-J.; Tkatchenko, A.; Müller, K.-R.
        Schnet - a Deep Learning Architecture for Molecules and Materials [J].
        The Journal of Chemical Physics, 2018, 148(24): 241722.

    Args:
        dim_feature (int, optional): Dimension of atomic representation. Defaults to 64.
        dim_filter (int, optional): Dimension of filter network. Defaults to 64.
        dim_edge_emb (int, optional): Dimension of edge embedding vectors. Defaults to None.
        interaction (Union[Interaction, List[Interaction], None], optional):
            Interaction module or list of interaction modules. Defaults to None.
        n_interaction (int, optional): Number of interaction layers. Defaults to 3.
        activation (Union[nn.Module, str], optional): Activation function. Defaults to 'ssp'.
        normalize_filter (bool, optional): Whether to normalize the filter network.
            Defaults to False.
        coupled_interaction (bool, optional): Whether to use coupled (shared) interaction layers.
            Defaults to False.
    
    Symbols:
        B: Number of simulation walker.
        A: Number of atoms in system.
        N: Number of neighbour atoms.
        D: Dimension of position coordinates, usually is 3.
        K: Number of basis functions in RBF.

    """

    def __init__(self,
                 dim_feature: int = 64,
                 dim_filter: int = 64,
                 dim_edge_emb: int = None,
                 interaction: Union[Interaction, List[Interaction]] = None,
                 n_interaction: int = 3,
                 activation: Union[nn.Module, str] = 'ssp',
                 normalize_filter: bool = False,
                 coupled_interaction: bool = False,
                 **kwargs,
                 ):

        super().__init__(
            dim_node_rep=dim_feature,
            dim_edge_rep=dim_feature,
            interaction=interaction,
            n_interaction=n_interaction,
            activation=activation,
            coupled_interaction=coupled_interaction,
            dim_node_emb=dim_feature,
            dim_edge_emb=dim_edge_emb,
        )
        self._kwargs = get_arguments(locals(), kwargs)

        self.dim_feature = get_integer(dim_feature)
        self.dim_filter = get_integer(dim_filter)
        self.normalize_filter = normalize_filter

        if self.interaction is None and self.dim_edge_emb is not None:
            self.build_interaction()

        self.default_embedding = self.get_default_embedding('schnet')

    def build_interaction(self):
        if self.dim_edge_emb is None:
            raise ValueError('Cannot build interaction without `dim_edge_emb`. '
                             'Please use `set_embedding_dimension` at first.')

        interaction_kwargs = {
            'dim_feature': self.dim_feature,
            'dim_edge_emb': self.dim_edge_emb,
            'dim_filter': self.dim_filter,
            'activation': self.activation,
            'normalize_filter': self.normalize_filter,
        }
        
        if self.coupled_interaction:
            self.interaction = nn.ModuleList([
                SchNetInteraction(**interaction_kwargs)
            ] * self.n_interaction)
        else:
            self.interaction = nn.ModuleList([
                SchNetInteraction(**interaction_kwargs)
                for _ in range(self.n_interaction)])